#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#define _USE_MATH_DEFINES
#include <math.h>

#include "raylib.h"
#include "pub_lib.h"
#include "FreeImage.h"
#include "image.h"
#include "dpi_algo.h"
#include "fftw3.h"


static void geo_mean_filter(unsigned int index, int* sum, unsigned char k,\
	int dim, BYTE val, BYTE* new_val, void* extra)
{
	int end = dim * dim;
	float *a = (float*)extra;
	if (k == 1)
		*a = (float)val / 255;
	else
		*a = (*a) * ((float)val / 255);

	float final = *a;
	if (k == end) {
		final = powf(final, 1.0f/end);
		final = final * 255;
		*new_val = (BYTE)final;
	}
}

//Geometric mean filter
void image_filter_geo_mean(FIBITMAP* ori, FIBITMAP* chg, int dim)
{
	float val;
	imgage_filter_common(ori, chg, dim, geo_mean_filter, (void*)&val);
}

/*
 * 5.3.3 local noise reduction filter
 * */
struct spatial_local_adaptive_data{
	unsigned char* data;
	float variance;
};

static void image_filter_local_adaptive(unsigned int index, int* sum,
	unsigned char k, int dim, BYTE val, BYTE* new_val, void* extra)
{
	int end = dim * dim;
	struct spatial_local_adaptive_data* p;
	p = (struct spatial_local_adaptive_data*)extra;
	p->data[k - 1] = val;

	if (k == end) {
		float sn = p->variance;
		float ml = data_get_avg(p->data, end);
		float sl = data_get_sigma2(p->data, end, ml);
		unsigned char gxy = p->data[end/2];
		float fxy;
		if (sn > sl)
			fxy = ml;
		else
			fxy = gxy - (sn / sl) * (gxy - ml);
		*new_val = (BYTE)fxy;
	}
}

void spatial_local_adaptive(FIBITMAP* ori, FIBITMAP* chg, int dim, float variance)
{
	struct spatial_local_adaptive_data p;
	p.variance = variance;
	p.data = malloc(dim * dim);
	imgage_filter_common(ori, chg, dim, image_filter_local_adaptive, (void*)&p);
	free(p.data);
}

/* Perform adaptive median filtering.
 * The median filter starts at size 3-by-3 and iterates
 * up to size SMAX-by-SMAX.
 * SMAX must be an odd integer greater than 1.
 * */
BYTE smedian_adaptive_filter(FIBITMAP* dib, int x, int y, int smax, BYTE* data)
{
	int p, q, dis, k, end;
	int ret, sxy;
	BYTE zmin, zmax, zmed, zxy;
	sxy = 3;
	FreeImage_GetPixelIndex(dib, x, y, &zxy);

stage_a:
	dis = sxy / 2;
	k = 0;
	end = sxy * sxy;

	for(q = y - dis; q <= y + dis; q++) {
		for(p = x - dis; p <= x + dis; p++) {
			FreeImage_GetPixelIndex(dib, p, q, &data[k]);
			k++;
		}
	}

	ret = kth_get_min(data, end);
	zmin = data[ret];
	ret = kth_get_max(data, end);
	zmax = data[ret];
	ret = kth_get_med(data, end);
	zmed = data[ret];

	if (zmed > zmin && zmed < zmax) {
		//stage b
		if(zxy > zmin && zxy < zmax)
			return zxy;
		else
			return zmed;
	} else {
		//increase the window size
		sxy++;
		if(sxy <= smax)
			goto stage_a;
		else
			return zmed;
	}

	return 0;
}

/*
 * 5.3.3 Adaptive, median filter
 * */
void spatial_median_adaptive(FIBITMAP* ori, FIBITMAP* chg, int smax)
{
	int x, y;
	BYTE new_value;

	int w, h;
	w = FreeImage_GetWidth(ori);
	h = FreeImage_GetHeight(ori);

	int dis = smax / 2;
	BYTE *data = malloc(smax * smax);
	for(y = dis; y < (h - dis); y++){
		for(x = dis; x < (w - dis); x++) {
			new_value = smedian_adaptive_filter(ori, x, y, smax, data);
			FreeImage_SetPixelIndex(chg, x, y, &new_value);
		}
	}
	free(data);
}

//example 5.10
fftwf_complex* freq_modeling_a(int w, int h, float a, float b, float T)
{
	int i, j, index, cplx;
	cplx = w/2 + 1;
	fftwf_complex *out = fftwf_alloc_complex(cplx * h);

	float alpha, sina, cosa, real, imag, k;
	for(j = 0; j < h; j++) {
		for(i = 0; i < cplx; i++) {
			index = j * cplx + i;
			//alpha = M_PI * (i * a + j * b);
			alpha = M_PI * ((i-w/2) * a + (j-h/2) * b);
			if (alpha < MIN_NUMBER)
				alpha = MIN_NUMBER;
			sina = sinf(alpha);
			cosa = cosf(alpha);
			k = sina * T / alpha;
			real = k * cosa;
			imag = 0 - k * sina;
			out[index] = build_complex(real, imag);
		}
	}

	return out;
}

//calc 1/in
void reciprocal_convert(fftwf_complex *in, int w, int h)
{
	int i;
	int size = (w/2 + 1) * h;
	fftwf_complex a;
	float real, imag, abs, k;

	for(i = 0; i < size; i++) {
		a = in[i];
		real = crealf(a);
		imag = cimagf(a);
		abs = cabsf(a);
		k = 1 / (abs * abs);
		real = k * real;
		imag = -k * imag;
		in[i] = build_complex(real, imag);
	}
}

//example 5.11
//calc 1/H
fftwf_complex* freq_degradation_a(int w, int h, float k, float r)
{
	int i, j, index, cplx;
	cplx = w/2 + 1;
	fftwf_complex *out = fftwf_alloc_complex(cplx * h);

	float a, b, temp, cx, cy;
	float scale, rr;
	cx = (float)w/2;
	cy = (float)h/2;

	rr = r * r;
	for(j = 0; j < h; j++) {
		for(i = 0; i < cplx; i++) {
			a = (i - cx) * (i - cx);
			b = (j - cy) * (j - cy);
			temp = k * powf((a+b), 5.0f/6);
			temp = expf(temp);

			scale = 1.0f;
			if(r > MIN_NUMBER) {
				//Buttworth, order = 10
				scale = (a + b) / rr;
				scale = powf(scale, 10);
				scale = 1.0f / (1.0f + scale);
			}
			temp = temp * scale;
			index = j * cplx + i;
			out[index] = build_complex(temp, 0);
		}
	}

	return out;
}

//example 5.12, Winer filter
//rH is reciprocal of H (1/H)
fftwf_complex* freq_degradation_wiener(int w, int h, fftwf_complex* rH, float K)
{
	int i, j, index, cplx;
	cplx = w/2 + 1;
	fftwf_complex *out = fftwf_alloc_complex(cplx * h);

	double hh, real, imag, rr, ii, res;
	fftwf_complex val;
	for(j = 0; j < h; j++) {
		for(i = 0; i < cplx; i++) {
			index = j * cplx + i;
			//calc |H|^2
			val = rH[index];
			real = crealf(val);
			imag = cimagf(val);
			rr = real * real;
			ii = imag * imag;
			if((rr + ii) < MIN_NUMBER)
				hh = 1 / MIN_NUMBER;
			else
				hh = 1 / (rr + ii);
			res = hh / (hh + K);

			real = real * res;
			imag = imag * res;
			out[index] = build_complex(real, imag);
		}
	}

	return out;
}


